package cmd

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/hex"
	"fmt"
	"os"

	"gitee.com/qianyucc/secure-tool/util"

	"github.com/spf13/cobra"
)

// aesCmd represents the aes command
var aesCmd = &cobra.Command{
	Use:   "aes",
	Short: "Encrypt and decrypt files and strings using AES algorithm",
	Long: `Encrypt and decrypt files and strings using AES algorithm.
	Use PKCS7 padding, CBC mode, and if you do not enter iv, the system will use key as iv`,
	Run: func(cmd *cobra.Command, args []string) {
		if err := cmd.ParseFlags(args); err != nil {
			fmt.Println(err)
			return
		}
		if Key == "" {
			fmt.Println("invalid parameters, please enter the secret key.")
			return
		}
		Iv = util.GetOrDefaultString(Iv, Key)

		if Text != "" {
			var raw []byte
			var err error
			if IsDecode {
				var textBytes []byte
				if textBytes, err = hex.DecodeString(Text); err != nil {
					fmt.Println(err)
					return
				}
				if raw, err = AesDecrypt(textBytes, []byte(Key), []byte(Iv)); err != nil {
					fmt.Println(err)
					return
				}
				util.PrintAndCopyTextToClipboardIgnoreError(string(raw))
			} else {
				if raw, err = AesEncrypt([]byte(Text), []byte(Key), []byte(Iv)); err != nil {
					fmt.Println(err)
					return
				}
				util.PrintAndCopyTextToClipboardIgnoreError(hex.EncodeToString(raw))
			}
			return
		}
		if InputFile != "" {
			if IsDecode {
				// 解密文件
				OutputFile = util.GetOrDefaultString(OutputFile, fmt.Sprintf("%s.dec", InputFile))
				if err := AesFileDecrypt(InputFile, OutputFile, []byte(Key), []byte(Iv)); err != nil {
					fmt.Println(err)
				}
			} else {
				OutputFile = util.GetOrDefaultString(OutputFile, fmt.Sprintf("%s.enc", InputFile))
				// 加密文件
				if err := AesFileEncrypt(InputFile, OutputFile, []byte(Key), []byte(Iv)); err != nil {
					fmt.Println(err)
				}
			}
			return
		}
		fmt.Println("invalid parameters, please enter the input file name or the string you want to encrypt or decrypt.")
	},
}

// PKCS7Padding PKCS7 填充
func PKCS7Padding(raw []byte, blockSize int) []byte {
	paddingInteger := blockSize - len(raw)%blockSize
	paddingBytes := bytes.Repeat([]byte{byte(paddingInteger)}, paddingInteger)
	return append(raw, paddingBytes...)
}

// PKCS7UnPadding 去除 PKCS7 填充
func PKCS7UnPadding(raw []byte) []byte {
	length := len(raw)
	paddingInteger := int(raw[length-1])
	return raw[:(length - paddingInteger)]
}

//AesEncrypt 加密函数
func AesEncrypt(raw []byte, key, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	blockSize := block.BlockSize()
	raw = PKCS7Padding(raw, blockSize)
	blockMode := cipher.NewCBCEncrypter(block, iv)
	result := make([]byte, len(raw))
	blockMode.CryptBlocks(result, raw)
	return result, nil
}

// AesDecrypt 解密函数
func AesDecrypt(raw []byte, key, iv []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	blockSize := block.BlockSize()
	blockMode := cipher.NewCBCDecrypter(block, iv[:blockSize])
	result := make([]byte, len(raw))
	blockMode.CryptBlocks(result, raw)
	result = PKCS7UnPadding(result)
	return result, nil
}

// AesFileEncrypt AES文件加密
func AesFileEncrypt(inputFile, outputFile string, key, iv []byte) (err error) {
	var fpInput, fpOutput *os.File
	if fpInput, err = os.Open(inputFile); err != nil {
		return err
	}
	if fpOutput, err = os.OpenFile(outputFile, os.O_WRONLY|os.O_CREATE, 0644); err != nil {
		return err
	}
	defer fpInput.Close()
	defer fpOutput.Close()

	var fileSize int64
	if fs, err := os.Stat(inputFile); err != nil {
		return err
	} else {
		fileSize = fs.Size()
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	blockSize := block.BlockSize()
	blockMode := cipher.NewCBCEncrypter(block, iv)
	sizeInOnceRead := blockSize * 10000
	// 加密次数
	var count int64
	isComplete := fileSize%int64(sizeInOnceRead) == 0
	if isComplete {
		count = fileSize / int64(sizeInOnceRead)
	} else {
		count = fileSize/int64(sizeInOnceRead) + 1
	}
	for i := 0; i < int(count); i++ {
		tempInputBytes := make([]byte, sizeInOnceRead)
		sz, _ := fpInput.Read(tempInputBytes)
		if i == int(count)-1 {
			tempInputBytes = PKCS7Padding(tempInputBytes[:sz], blockSize)
		}
		tempResult := make([]byte, len(tempInputBytes))
		blockMode.CryptBlocks(tempResult, tempInputBytes)
		_, err = fpOutput.Write(tempResult)
		if err != nil {
			return err
		}
	}
	return nil
}

// AesFileDecrypt AES文件解密
func AesFileDecrypt(inputFile, outputFile string, key, iv []byte) (err error) {
	var fpInput, fpOutput *os.File
	if fpInput, err = os.Open(inputFile); err != nil {
		return err
	}
	if fpOutput, err = os.OpenFile(outputFile, os.O_WRONLY|os.O_CREATE, 0644); err != nil {
		return err
	}
	defer fpInput.Close()
	defer fpOutput.Close()

	var fileSize int64
	if fs, err := os.Stat(inputFile); err != nil {
		return err
	} else {
		fileSize = fs.Size()
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return err
	}
	blockSize := block.BlockSize()
	blockMode := cipher.NewCBCDecrypter(block, iv)
	sizeInOnceRead := blockSize * 10000
	var count int64
	isComplete := fileSize%int64(sizeInOnceRead) == 0
	if isComplete {
		count = fileSize / int64(sizeInOnceRead)
	} else {
		count = fileSize/int64(sizeInOnceRead) + 1
	}
	for i := 0; i < int(count); i++ {
		tempInputBytes := make([]byte, sizeInOnceRead)
		sz, _ := fpInput.Read(tempInputBytes)
		tempResult := make([]byte, sz)
		blockMode.CryptBlocks(tempResult, tempInputBytes[:sz])
		if i == int(count)-1 {
			tempResult = PKCS7UnPadding(tempResult)
		}
		if _, err = fpOutput.Write(tempResult); err != nil {
			return err
		}
	}
	return nil
}

func init() {
	rootCmd.AddCommand(aesCmd)

	aesCmd.Flags().StringVarP(&InputFile, "file", "f", "", "The file name you want to encrypt or decrypt")
	aesCmd.Flags().StringVarP(&OutputFile, "out", "o", "", `File name after encryption or decryption. 
	When encrypting, if you do not enter this flag, the output file will be named as ${fileName}.enc by default,
	When decrypting, if you do not enter this flag, the output file will be named as ${fileName}.dec by default`)
	aesCmd.Flags().BoolVarP(&IsDecode, "decode", "d", false, "Indicates whether it is a decryption operation")
	aesCmd.Flags().StringVarP(&Key, "key", "k", "", "Encryption or decryption key")
	aesCmd.Flags().StringVarP(&Iv, "iv", "i", "", "Encryption or decryption iv, if you do not enter iv, the system will use key as iv")
	aesCmd.Flags().StringVarP(&Text, "text", "t", "", "The string you want to encrypt or decrypt")
}
